#!/usr/bin/env python
'''
Class representing Lattice
Created on Jan 29, 2011

@author: vladimir
'''
from __future__ import division
import numpy
import random
import math
import sys

def xcombinations(m):
    '''
    generates all tensor product for our lattice
    '''
    for i in outerRange2(len(m), m):
        yield i
    
def outerRange2(d, m):
    if d > 0:
        for i in range(0, m[-d]):
            for l in outerRange2(d - 1, m):
                yield [i] + l
    else:
        yield []

class Lattice:
    '''
    class represents lattice of interest
    '''
    def __init__(self, h_shape, v_size, h_neighbours, ferromagnetic, **kwargs ):
        
        self.h_size = h_shape[0]
        self.v_size = v_size
        self.shape = tuple( [v_size] + list(h_shape) )
        self.ferromagnetic = ferromagnetic
        if ferromagnetic:
            self.J = 1
        else:
            self.J = -1
            
        self.lattice = numpy.zeros( self.shape )
        self.LD = kwargs.get('LD', 0)
        self.T = kwargs['T']
        self.B = kwargs['B']
        
        self.h_neighbours = {}
        
        v_neighbors = kwargs.get('v_neighbours', ())
        self.v_neighbors = v_neighbors
        self.v_neighbours = {}
        for index in self:
            self.lattice[ index ] = random.choice([-1, 1])
            th = []
            tv = []
            for neighbor in h_neighbours:
                tn = numpy.array(  [ 0 ] + list(neighbor) ) + index
                for i in xrange(1, len(tn)):
                    if tn[i] >= self.h_size:
                        tn[i] = 0
                    elif tn[i] == -1:
                        tn[i] = self.h_size - 1
                th += [ tuple(tn) ]
            self.h_neighbours[ index ] = tuple(th)
            
            if self.v_size > 1:    
                for neighbor in v_neighbors:
                    tn = numpy.array( [neighbor] + len(h_shape) * [0] ) + index
                    if tn[0] >= self.v_size:
                        tn[0] = 0
                    elif tn[0] == -1:
                        tn[0] = self.v_size - 1
                    tv += [ tuple(tn) ]
                self.v_neighbours[ index ] = tuple(tv)

    def Energy_J(self):
        '''
        calculates J part of Energy of the lattice
        '''
        result = 0
        for index in self:
            result += self.Sum_h_sites( index ) * self.lattice[ index ] 
                
        return - self.J * result / 2
    
    def Energy_B(self):
        '''
        calculates B part of Energy of the lattice
        '''
        if self.B == 0:
            return 0
        elif not self.v_neighbors:
            return 0
        elif self.v_size == 1:
            return 0
        else:
            result = 0
            for index in self:
                result += math.exp( -2 * self.LD * self.lattice[ index ] * self.lattice[ self.v_neighbours[ index][0] ] )
            return - self.B * result
        
    def Sum_h_sites(self, position):
        '''
        calculates sum of the spins of nearest horizontal neighbors near site on position - position
        '''
        result = 0        
        for neighbor in self.h_neighbours[ position ]:
            result += self.lattice[ neighbor ]
        return result
    
    def Sum_v_sites(self, position):
        '''
        calculates sum of the spins of nearest vertical neighbors near site on position - position
        '''
        if self.v_size > 1:
            result = 0
            for neighbor in self.v_neighbours[ position ]:
                result += self.lattice[ neighbor ]
            return result
        else:
            return 0
    
    def __iter__(self):
        '''
        iterator over all sites
        '''
        if 'viter' not in dir(self):
            self.viter = [ tuple(i) for i in xcombinations( self.shape ) ]
            
        for i in self.viter:
            yield i
            
    
    def __str__(self):
        '''
        what to print
        '''
        return str(self.lattice)
    
    def sum_layer(self, L):
        '''
        returns sum of the spin in the L'th layer 
        '''
        
        return sum( [numpy.sum( self.lattice[i]) for i in xrange( self.v_size)] ) / self.v_size
        
    def RandomSite(self):
        '''
        Returns random site of the lattice
        '''
        result = []
        for i in self.shape:
            result += [ random.randint(0,  i - 1) ]
        return tuple(result)
